import torch
import torch.nn as nn

def myCE(x1, x2):
    # x1, x2 [batch_size, num_classes]
    m = nn.Softmax(dim=1)
    x1 = m(x1)
    x2 = m(x2)
    return -torch.sum(x2 * x1.log(), dim=1)

def CE_from_label(outputs, labels):
    labels = torch.unsqueeze(labels, dim=-1)
    labels = torch.zeros_like(outputs).scatter_(1, labels, 1)
    return myCE(outputs, labels)
    
def filter(outputs, labels, noise_rate):
    labels = torch.unsqueeze(labels, dim=-1)
    one_hot = torch.zeros_like(outputs).scatter_(1, labels, 1)
    loss = myCE(outputs, one_hot)
    kth = int(outputs.size()[0] * (1 - noise_rate))
    loss, _ = torch.topk(loss, kth, largest=False)
    return loss.mean()

def filter_replace(outputs1, outputs2, labels, noise_rate):
    labels = torch.unsqueeze(labels, dim=-1)
    labels = torch.zeros_like(outputs1).scatter_(1, labels, 1)
    loss0 = myCE(outputs1, outputs2.detach())
    loss1 = myCE(outputs1, labels)
    loss2 = myCE(outputs2, labels)
    kth = int((1 - noise_rate) * labels.shape[0])
    tao, pos = torch.topk(loss1, kth, largest=False)
    # _, pos = torch.topk(loss2, kth, largest=False)
    student_loss = torch.where(loss1 < tao[-1], loss1, loss0)
    teacher_loss = torch.take(loss2, pos)
    return student_loss.mean(), teacher_loss.mean()

def filter_by_tao(outputs, labels, tao):
    labels = torch.unsqueeze(labels, dim=-1)
    one_hot = torch.zeros_like(outputs).scatter_(1, labels, 1)
    loss = myCE(outputs, one_hot)
    loss = torch.masked_select(loss, loss < tao)
    return loss.mean()

def filter_replace_by_tao(outputs1, outputs2, labels, tao):
    labels = torch.unsqueeze(labels, dim=-1)
    labels = torch.zeros_like(outputs1).scatter_(1, labels, 1)
    loss0 = myCE(outputs1, outputs2.detach())
    loss1 = myCE(outputs1, labels)
    loss2 = myCE(outputs2, labels)
    student_loss = torch.where(loss1 < tao, loss1, loss0)
    teacher_loss = torch.masked_select(loss2, loss1 < tao)
    return student_loss.mean(), teacher_loss.mean()


if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


    x_train = torch.tensor([[3.3, 4.4], [5.5, 6.71], [6.93, 4.168], [9.779, 6.182], [7.59, 2.167], [7.042, 10.791], [5.313, 7.997]], dtype=torch.float).to(device)

    y_train = torch.tensor([[1.7, 2.76], [2.09, 3.19], [1.694, 1.573], [3.366, 2.596], [2.53, 1.221], [2.827, 3.465], [1.65, 2.904]], dtype=torch.float).to(device)

    x_train.requires_grad = True
    y_train.requires_grad = True


    l = torch.tensor([1, 1, 0, 1, 0, 0, 1]).to(device)

    filter_by_tao(x_train, l, 1)
    input()
    print(filter(x_train, l, 0.2))
    print(filter_replace(x_train, y_train, l, 0.2))

    # 183417653